package readers

import (
	"context"
	"crypto/md5"
	"database/sql"
	"fmt"
	"log"
	"mig2kwdb/common"
	"strconv"
	"sync"
	"time"

	"github.com/jackc/pglogrepl"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgproto3"
)

type PostgresqlReplicator struct {
	conn       *pgconn.PgConn
	binIndex   int
	mutex      sync.Mutex
	ctx        context.Context
	handler    common.ReplicateHandlerInterface
	err        error
	waitgGroup sync.WaitGroup
	database   common.DatabaseConfig
	dbconInfo  common.DBConnectionInfo
	logPos     pglogrepl.LSN
	tables     []string
}

func (t *PostgresqlReplicator) GetPublicationName() string {
	md5code := fmt.Sprintf("%x", md5.Sum([]byte(t.database.Name)))
	publication_name := `pub_pub_` + md5code
	return publication_name
}
func (t *PostgresqlReplicator) GetReplicationSlotName() string {
	md5code := fmt.Sprintf("%x", md5.Sum([]byte(t.database.Name)))
	publication_name := `slot_rep_` + md5code
	return publication_name
}
func (t *PostgresqlReplicator) Stop() {
	publication_name := t.GetPublicationName()
	replication_slot_name := t.GetReplicationSlotName()
	if t.conn == nil {
		return
	}
	//set the identify to default
	sql_str := ``
	for _, table := range t.tables {
		sql_str += `ALTER TABLE ` + t.database.Schema + "." + table + ` REPLICA IDENTITY default;`
	}
	t.waitgGroup.Wait()

	t.mutex.Lock()
	defer t.mutex.Unlock()

	result := t.conn.Exec(context.Background(), `DROP PUBLICATION IF EXISTS `+publication_name+`;`)
	result.ReadAll()
	if t.conn != nil {
		t.conn.Close(context.Background())
	}
	t.conn = nil
	//drop slot
	//must create a new conn to drop the replication slot
	dsn := "postgres://" + t.dbconInfo.User + ":" + t.dbconInfo.Password + "@" + t.dbconInfo.Ip + ":" + fmt.Sprint(t.dbconInfo.Port) + "/" + t.database.Name + "?replication=database"

	rep_conn, rep_err := pgconn.Connect(context.Background(), dsn)
	if rep_err == nil {
		dropOption := pglogrepl.DropReplicationSlotOptions{Wait: true}
		pglogrepl.DropReplicationSlot(context.Background(), rep_conn, replication_slot_name, dropOption)
		rep_conn.Close(context.Background())
	}

	dsn = "postgres://" + t.dbconInfo.User + ":" + t.dbconInfo.Password + "@" + t.dbconInfo.Ip + ":" + fmt.Sprint(t.dbconInfo.Port) + "/" + t.database.Name + "?"

	if t.dbconInfo.SSL {
		dsn += "sslmode=require"
	} else {
		dsn += "sslmode=disable"
	}

	dbconn, err := sql.Open("pgx", dsn)

	if err == nil {
		//maybe execute long time
		ctx, cancelFunc := context.WithTimeout(context.Background(), time.Second*3)
		go func() {
			dbconn.ExecContext(ctx, sql_str)
			cancelFunc()
		}()
		<-ctx.Done()
		dbconn.Close()
		dbconn = nil
	}

	//log.Println("replicate stopped")
}
func (t *PostgresqlReplicator) GetDsnByConninfo(dbconInfo common.DBConnectionInfo, database common.DatabaseConfig) string {
	dsn := "postgres://" + dbconInfo.User + ":" + dbconInfo.Password + "@" + dbconInfo.Ip + ":" + fmt.Sprint(dbconInfo.Port) + "/" + database.Name + "?replication=database"
	return dsn
}
func (t *PostgresqlReplicator) Init(ctx context.Context, handler common.ReplicateHandlerInterface,
	database common.DatabaseConfig, dbconInfo common.DBConnectionInfo, tables []string) error {
	t.handler = handler
	t.ctx = ctx
	t.err = nil
	t.dbconInfo = dbconInfo
	t.database = database
	t.tables = tables

	dsn := t.GetDsnByConninfo(t.dbconInfo, t.database)
	var err error = nil
	if t.conn == nil {
		t.conn, err = pgconn.Connect(context.Background(), dsn)
		if err != nil {
			log.Println(err)
			return err
		}
	}
	const outputPlugin = "pgoutput"
	publication_name := t.GetPublicationName()
	replication_slot_name := t.GetReplicationSlotName()
	for _, table := range t.tables {
		sql_str := `ALTER TABLE ` + database.Schema + "." + table + ` REPLICA IDENTITY FULL;`
		result := t.conn.Exec(context.Background(), sql_str)
		_, err := result.ReadAll()
		if err != nil {
			log.Println(err)
			return err
		}
	}
	result := t.conn.Exec(context.Background(), `DROP PUBLICATION IF EXISTS `+publication_name+`;`)
	_, err = result.ReadAll()
	if err != nil {
		log.Println("drop publication if exists error", err)
		return err
	}
	dropOption := pglogrepl.DropReplicationSlotOptions{Wait: false}
	pglogrepl.DropReplicationSlot(context.Background(), t.conn, replication_slot_name, dropOption)

	//create the publication

	for_tables := ` FOR ALL TABLES; `
	if len(t.tables) > 0 {
		for_tables = ` FOR TABLE `
		for idx, tab := range t.tables {
			if idx > 0 {
				for_tables += `,`
			}
			for_tables += t.database.Schema + "." + tab
		}
	}
	create_pub := `CREATE PUBLICATION ` + publication_name + for_tables
	result = t.conn.Exec(context.Background(), create_pub)
	_, err = result.ReadAll()
	if err != nil {
		log.Println(err, create_pub)
		return err
	}
	_, err = pglogrepl.CreateReplicationSlot(context.Background(), t.conn, replication_slot_name, outputPlugin, pglogrepl.CreateReplicationSlotOptions{Temporary: false})
	if err != nil {
		log.Println("CreateReplicationSlot failed:", err)
		return err
	}
	sysident, err := pglogrepl.IdentifySystem(context.Background(), t.conn)
	if err != nil {
		log.Println(err)
		return err
	}
	//log.Println("SystemID:", sysident.SystemID, "Timeline:", sysident.Timeline, "LogPos:", sysident.XLogPos, "DBName:", sysident.DBName)
	t.logPos = sysident.XLogPos
	return nil
}
func (t *PostgresqlReplicator) Start(ctx context.Context) error {

	t.Replicate(ctx)
	return nil
}

// getColValues 将逻辑复制消息中的列数据解析为 *sql.NullString 列表
func (t *PostgresqlReplicator) getColValues(columns []*pglogrepl.TupleDataColumn, cols []common.KwdbColumn) []interface{} {
	values := make([]interface{}, len(columns))
	for i := range columns {
		values[i] = &sql.NullString{}
	}

	for idx, col := range columns {
		switch col.DataType {
		case 'n', 'u': // null
			values[idx] = &sql.NullString{}
		case 't': // text
			values[idx] = &sql.NullString{
				String: string(col.Data),
				Valid:  true,
			}

		}
	}
	return values
}

// getColumns 根据 RelationID 获取 column 信息，如果不存在则记录日志并返回 nil
func (t *PostgresqlReplicator) getColumns(tablesMap map[string][]common.KwdbColumn, relations map[uint32]*pglogrepl.RelationMessage, relationID uint32) (string, []common.KwdbColumn) {
	rel, ok := relations[relationID]
	if !ok {
		log.Println("unknown relation ID ", relationID)
		return "", nil
	}
	dbtb := t.database.Schema + "." + rel.RelationName
	cols, ok := tablesMap[dbtb]
	if !ok || len(cols) == 0 {
		columns, _ := t.handler.GetTableColumn(t.database.Name, t.database.Schema, rel.RelationName)
		tablesMap[dbtb] = columns
		cols = columns
	}
	return rel.RelationName, cols
}
func (t *PostgresqlReplicator) Replicate(ctx context.Context) {
	t.waitgGroup.Add(1)
	defer t.Stop()
	defer t.waitgGroup.Done()

	replication_slot_name := t.GetReplicationSlotName()
	var pluginArguments = []string{"proto_version '1'", `publication_names '` + t.GetPublicationName() + `'`}
	sysident, err := pglogrepl.IdentifySystem(context.Background(), t.conn)
	if err != nil {
		log.Println(err)
		return
	}
	position, _ := strconv.ParseUint(strconv.FormatUint(uint64(sysident.XLogPos), 10), 10, 64)
	err = pglogrepl.StartReplication(context.Background(), t.conn, replication_slot_name, pglogrepl.LSN(position), pglogrepl.StartReplicationOptions{PluginArgs: pluginArguments})
	if err != nil {
		log.Println("StartReplication failed:", err)
		return
	}

	standbyMessageTimeout := time.Second * 10
	nextStandbyMessageDeadline := time.Now().Add(standbyMessageTimeout)
	relations := map[uint32]*pglogrepl.RelationMessage{}
	tablesMap := make(map[string][]common.KwdbColumn)
	for {
		if common.IsStopped(ctx) {
			break
		}
		if time.Now().After(nextStandbyMessageDeadline) {
			err = pglogrepl.SendStandbyStatusUpdate(context.Background(), t.conn, pglogrepl.StandbyStatusUpdate{WALWritePosition: t.logPos})
			if err != nil {
				log.Fatalln("SendStandbyStatusUpdate failed:", err)
			}
			nextStandbyMessageDeadline = time.Now().Add(standbyMessageTimeout)
		}
		ctx, cancel := context.WithDeadline(context.Background(), nextStandbyMessageDeadline)
		rawMsg, err := t.conn.ReceiveMessage(ctx)
		cancel()
		if err != nil {
			if pgconn.Timeout(err) {
				time.Sleep(time.Second)
				continue
			} else {
				log.Println("ReceiveMessage failed:", err)
				return
			}
		}
		if errMsg, ok := rawMsg.(*pgproto3.ErrorResponse); ok {
			log.Println(errMsg)
			return
		}

		msg, ok := rawMsg.(*pgproto3.CopyData)
		if !ok {
			continue
		}

		switch msg.Data[0] {
		case pglogrepl.PrimaryKeepaliveMessageByteID:
			pkm, err := pglogrepl.ParsePrimaryKeepaliveMessage(msg.Data[1:])
			if err != nil {
				log.Println("ParsePrimaryKeepaliveMessage failed:", err)
			}
			if pkm.ServerWALEnd > t.logPos {
				t.logPos = pkm.ServerWALEnd
			}
			if pkm.ReplyRequested {
				nextStandbyMessageDeadline = time.Time{}
			}

		case pglogrepl.XLogDataByteID:
			xld, err := pglogrepl.ParseXLogData(msg.Data[1:])
			if err != nil {
				log.Fatalln("ParseXLogData failed:", err)
			}
			//log.Println("XLogData =>", "WALStart", xld.WALStart, "ServerWALEnd", xld.ServerWALEnd, "ServerTime:", xld.ServerTime, "WALData", string(xld.WALData))
			logicalMsg, err := pglogrepl.Parse(xld.WALData)
			if err != nil {
				log.Fatalf("Parse logical replication message: %s", err)
			}
			//log.Printf("Receive a logical replication message: %s", logicalMsg.Type())
			switch logicalMsg := logicalMsg.(type) {
			case *pglogrepl.RelationMessage:
				relations[logicalMsg.RelationID] = logicalMsg

			case *pglogrepl.InsertMessage:
				tableName, cols := t.getColumns(tablesMap, relations, logicalMsg.RelationID)
				if cols == nil {
					continue
				}
				value := t.getColValues(logicalMsg.Tuple.Columns, cols)
				values := [][]interface{}{value}
				_, err := t.handler.ReplicateRecords(t.database.Name, t.database.Schema, tableName, cols, values, "insert")
				if err != nil {
					log.Println(err)
				}

			case *pglogrepl.UpdateMessage:
				tableName, cols := t.getColumns(tablesMap, relations, logicalMsg.RelationID)
				if cols == nil {
					continue
				}

				oldValues := t.getColValues(logicalMsg.OldTuple.Columns, cols)
				newValues := t.getColValues(logicalMsg.NewTuple.Columns, cols)
				values := [][]interface{}{oldValues, newValues}
				//log.Println(oldValues, newValues)
				_, err := t.handler.ReplicateRecords(t.database.Name, t.database.Schema, tableName, cols, values, "update")
				if err != nil {
					log.Println(err)
				}

			case *pglogrepl.DeleteMessage:
				tableName, cols := t.getColumns(tablesMap, relations, logicalMsg.RelationID)
				if cols == nil {
					continue
				}

				deleteValues := t.getColValues(logicalMsg.OldTuple.Columns, cols)
				values := [][]interface{}{deleteValues}
				_, err := t.handler.ReplicateRecords(t.database.Name, t.database.Schema, tableName, cols, values, "delete")
				if err != nil {
					log.Println(err)
				}

			}
			if xld.WALStart > t.logPos {
				t.logPos = xld.WALStart
			}
		}
	}

}
